# infer.py
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import json

# ------------------------------
# 1. 配置 4bit 量化
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4"
)

# ------------------------------
# 2. 加载基础模型
base_model_name = r"D:\models\Qwen3\Qwen3-1.7B"
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

# ------------------------------
# 3. 加载 LoRA 权重
lora_path = "../../qwen3-lora"
model = PeftModel.from_pretrained(base_model, lora_path)
model.eval()  # 推理模式

# ------------------------------
# 4. 加载分词器
tokenizer = AutoTokenizer.from_pretrained(base_model_name)


# ------------------------------
# 5. 推理函数
def generate_resume_skills(prompt, max_tokens=700):
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        do_sample=True,
        top_k=50,
        top_p=0.9,
        temperature=0.7
    )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # 尝试只提取 JSON 内容
    try:
        # 找到第一个 { 开头和最后一个 } 结尾
        start = text.index("{")
        end = text.rindex("}") + 1
        json_text = text[start:end]
        return json.loads(json_text)
    except Exception:
        return text  # 若解析失败，直接返回原文本


# ------------------------------
# 6. 测试推理
prompt = """你是一名简历优化专家，请根据以下输入生成优化后的技能描述，只需补充 output 字段内容。

{
  "instruction": "简历专业技能优化专家，专注于系统化梳理与优化技能描述。",
  "section": "技能",
  "job_target": "java开发工程师",
  "input": "熟练掌握JavaSe核心知识，对集合、多线程有深入的理解，了解JVM的内存模型、垃圾回收(GC)等知识。熟练掌握MySQL、达梦数据库，熟悉MySQL的事务、索引、锁，了解MVCC底层工作原理和SQL优化。熟悉Redis等非关系型数据库，了解Redis缓存机制，理解持久化机制和缓存相关问题。熟练掌握Spring、SpringMVC、SpingBoot、Mybatis、Mybatis-plus开发框架，熟知IOC、AOP、事务底层工作原理。熟知SpringCloud系列微服务相关概念以及Alibaba系列微服务组件。掌握Linux常用命令，对项目打包、配置、部署均有线上使用经验。熟知常用的设计模式，例如单例、模板、策略、代理、工厂等有线上使用经验。熟练使用开发工具Git、itrust堡垒机、Apifox、PDManer、禅道等辅助性工具，熟练使用OkHttp、Hutool、EasyExcel等第三方工具。了解ES、RabbitMQ、docker、vueSpringAI、大模型等。",
  "output":
}
"""

result = generate_resume_skills(prompt)
print(json.dumps(result, ensure_ascii=False, indent=2))
